/*
This is part of Screenbar, which is based in part on the Litestep
shell source code.

Copyright (C) 2008-2009 Jim Babcock
Copyright (C) 1997-2009 The LiteStep Development Team

This program is free software; you can redistribute it and/or
modify it under the terms of the GNU General Public License
as published by the Free Software Foundation; either version 2
of the License, or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; if not, write to the Free Software
Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
*/
#include "stdafx.h"
#include <boost/bind.hpp>

SnapshotManager *snapshotManager = NULL;

/////////////////////////////////////////////////////////////////////////////


SnapshotManager::SnapshotManager()
{
	HINSTANCE instUser32 = LoadLibraryA("user32.dll");
	printWindowFunc = (PrintWindowType)GetProcAddress(instUser32, "PrintWindow");
	
	setConfigBool("CanSnapshot", printWindowFunc != NULL);
	
	exiting = false;
	snapper = new thread(boost::bind(&SnapshotManager::snapshotLoop, this));
}

SnapshotManager::~SnapshotManager()
{
	sync.lock();
		for(unsigned ii=0; ii<pending.size(); ii++)
			pending[ii]->request->updateCancelled = true;
		requests.clear();
		exiting = true;
	sync.unlock();
	requestsWaiting.notify_all();
	
	if(snapper)
	{
		snapper->join();
		delete snapper;
	}
}

SnapshotRequest *SnapshotManager::requestSnapshot(SnapshotSpec spec, HWND requestSource)
{
	SnapshotRequest *ret;
	
	sync.lock();
		map<SnapshotSpec, SnapshotRequest*>::iterator it = requests.find(spec);
		if(it != requests.end()) {
			ret = it->second;
			ret->addNotifyRequest(requestSource);
		} else {
			ret = new SnapshotRequest(spec);
			ret->addNotifyRequest(requestSource);
			requests[spec] = ret;
			requestUpdate(ret);
		}
		ret->incrRefCount();
	sync.unlock();
	return ret;
}

void SnapshotManager::cancelRequest(SnapshotRequest *request)
{
	sync.lock();
		requests.erase(request->getSpec());
		delete request;
	sync.unlock();
}

void SnapshotManager::updateSnapshots()
{
	UINT64 time = getSystemTime();
	
	UINT64 foregroundStaleTime = time - (UINT64)10000 * settings->foregroundSnapshotInterval;
	UINT64 backgroundStaleTime = time - (UINT64)10000 * settings->backgroundSnapshotInterval;
	
	sync.lock();
	
	for(map<SnapshotSpec, SnapshotRequest*>::iterator ii = requests.begin();
	    ii != requests.end(); ii++)
	{
		SnapshotRequest *request = ii->second;
		
		if(request->updatePending)
			continue;
		
		UINT64 timeTaken = request->getSnapshotTime();
		
		// Don't refresh a screenshot that's newer than the update period
		if(windowTracker->getForegroundTask() == ii->first.task) {
			if(timeTaken > foregroundStaleTime)
				continue;
		} else {
			if(timeTaken > backgroundStaleTime)
				continue;
		}
		
		requestUpdate(ii->second);
	}
	sync.unlock();
}

SnapshotLayout *SnapshotManager::getTaskSnapLayout(SnapshotRequest *request)
{
	SnapshotLayout *ret = new SnapshotLayout();
	SnapshotSpec spec = request->getSpec();
	
	// Enumerate this task's windows
	vector<WindowData*> windowData;
	spec.task->getWindows(windowData);
	int numWindows = windowData.size();
	
	for(int ii=0; ii<numWindows; ii++)
		ret->windows.push_back(windowData[ii]->handle);
	
	// Get the positions of all of this task's windows
	for(int ii=0; ii<numWindows; ii++) {
		RECT rect;
		GetWindowRect(ret->windows[ii], &rect);
		ret->windowRects.push_back(rect);
	}
	
	// Decide where in the output image the windows will be placed
	
	ret->targetRects = ret->windowRects;
	// Translate so that the to-leftmost spot is at (0,0)
	int leftBound = ret->targetRects[0].left;
	int topBound = ret->targetRects[0].top;
	for(int ii=1; ii<numWindows; ii++)
	{
		if(ret->targetRects[ii].left < leftBound)
			leftBound = ret->targetRects[ii].left;
		if(ret->targetRects[ii].top < topBound)
			topBound = ret->targetRects[ii].top;
	}
	for(int ii=0; ii<numWindows; ii++)
	{
		ret->targetRects[ii].left   -= leftBound;
		ret->targetRects[ii].right  -= leftBound;
		ret->targetRects[ii].top    -= topBound;
		ret->targetRects[ii].bottom -= topBound;
	}
	
	// Remove any empty columns or rows in the positioning, so that
	// if a window and its toolbox are on opposite ends of the screen, we
	// don't have to scale them down to show the gap between them.
	// TODO
	
	// Uniformly scale everything to the output size
	int rightBound = ret->targetRects[0].right;
	int botBound = ret->targetRects[0].bottom;
	for(int ii=1; ii<numWindows; ii++)
	{
		if(ret->targetRects[ii].right > rightBound)
			rightBound = ret->targetRects[ii].right;
		if(ret->targetRects[ii].bottom > botBound)
			botBound = ret->targetRects[ii].bottom;
	}
	
	float scaleFactor = 1.0;
	float thickness;
	float length;
	if(spec.vertical) {
		thickness = rightBound;
		length = botBound;
	} else {
		thickness = botBound;
		length = rightBound;
	}
	
	if(thickness > spec.thickness) {
		scaleFactor = (float)spec.thickness / thickness;
		thickness *= scaleFactor;
		length *= scaleFactor;
	}
	if(length > spec.maxLength) {
		float lengthAdjustment = spec.maxLength / length;
		scaleFactor *= lengthAdjustment;
		thickness *= lengthAdjustment;
		length *= lengthAdjustment;
	}
	
	ret->outputWidth = rightBound*scaleFactor;
	ret->outputHeight = botBound*scaleFactor;
	
	for(int ii=0; ii<numWindows; ii++)
	{
		ret->targetRects[ii].left   *= scaleFactor;
		ret->targetRects[ii].right  *= scaleFactor;
		ret->targetRects[ii].top    *= scaleFactor;
		ret->targetRects[ii].bottom *= scaleFactor;
	}
	
	
	return ret;
}

void SnapshotManager::requestUpdate(SnapshotRequest *request)
{
	request->incrRefCount();
	request->updatePending = true;
	
	PendingUpdate *pendingUpdate = new PendingUpdate();
	pendingUpdate->layout = getTaskSnapLayout(request);
	pendingUpdate->request = request;
	
	pending.push_back(pendingUpdate);
	requestsWaiting.notify_all();
}

/////////////////////////////////////////////////////////////////////////////

void SnapshotManager::snapshotLoop(SnapshotManager *snapshotManager)
{
	mutex waitMutex;
	boost::unique_lock<mutex> waitLock(waitMutex);
	
	while(!snapshotManager->exiting)
	{
		PendingUpdate *update = snapshotManager->getNextRequest();
		
		if(update)
		{
			SnapshotRequest *request = update->request;
			
			if(!request->updateCancelled)
			{
				Snapshot *snapshot = snapshotManager->snap(update->layout);
				
				request->setSnapshot(snapshot);
				request->updatePending = false;
				
				request->decrRefCount();
			}
			
			delete update->layout;
			delete update;
		} else {
			snapshotManager->requestsWaiting.wait(waitLock);
		}
	}
}

SnapshotManager::PendingUpdate *SnapshotManager::getNextRequest()
{
	PendingUpdate *ret;
	sync.lock();
	
	// TODO: Check that the windows are responding first
	if(pending.size()) {
		ret = pending[0];
		pending.pop_front();
	} else {
		ret = NULL;
	}
	
	sync.unlock();
	return ret;
}

Snapshot *SnapshotManager::snap(SnapshotLayout *layout)
{
	if(!printWindowFunc)
		return NULL;
	
	int outputWidth = layout->outputWidth;
	int outputHeight = layout->outputHeight;
	vector<HWND> &windows = layout->windows;
	vector<RECT> &windowRects = layout->windowRects;
	vector<RECT> &targetRects = layout->targetRects;
	
	if(!windows.size())
		return NULL;
	
	// Create a surface to draw onto
	HDC formatDC = GetDC(windows[0]);
	
	HDC combinedImage = CreateCompatibleDC(formatDC);
	HBITMAP combinedBmp = CreateCompatibleBitmap(formatDC, outputWidth, outputHeight);
	HBITMAP oldBmp = (HBITMAP)SelectObject(combinedImage, combinedBmp);

	ReleaseDC(windows[0], formatDC);

	// Fill it with magic pink
	HBRUSH fillBrush = CreateSolidBrush(RGB(255,0,255));
	RECT fillRect = {0, 0, outputWidth, outputHeight};
	FillRect(combinedImage, &fillRect, fillBrush);
	DeleteObject(fillBrush);
	
	XFORM identityTransform;
		identityTransform.eM11 = identityTransform.eM22 = 1.0;
		identityTransform.eM12 = identityTransform.eM21 = 0.0;
		identityTransform.eDx = identityTransform.eDy = 0.0;
	
	// Have windows paint themselves
	for(unsigned ii=0; ii<windows.size(); ii++)
	{
		HWND hwnd = windows[ii];
		
		// Set up the transform for painting
		RECT targetRect = targetRects[ii];
		int targetWidth = targetRect.right - targetRect.left;
		int targetHeight = targetRect.bottom - targetRect.top;
		int sourceWidth = windowRects[ii].right - windowRects[ii].left;
		int sourceHeight = windowRects[ii].bottom - windowRects[ii].top;
		
		XFORM transform;
		transform.eM12 = transform.eM21 = 0;
		transform.eM11 = (float)targetWidth / (float)sourceWidth;
		transform.eM22 = (float)targetHeight / (float)sourceHeight;
		transform.eDx = 0;
		transform.eDy = 0;
		
		// Allocate a bitmap to draw onto
		HDC windowDrawContext = CreateCompatibleDC(combinedImage);
		HBITMAP windowDrawBmp = CreateCompatibleBitmap(combinedImage,
			targetRect.right-targetRect.left, targetRect.bottom-targetRect.top);
		HBITMAP oldWindowBmp = (HBITMAP)SelectObject(windowDrawContext, windowDrawBmp);
		
		// Set painting options and paint
		/*HRGN region = getPreviewRegion(hwnd, &transform);
		HRGN oldRegion = NULL;
		if(region)
			oldRegion = (HRGN)SelectObject(windowDrawContext, region);*/
		
		SetGraphicsMode(windowDrawContext, GM_ADVANCED);
		SetMapMode(windowDrawContext, MM_ANISOTROPIC);
		SetStretchBltMode(windowDrawContext, HALFTONE);
		
		SetWorldTransform(windowDrawContext, &transform);
			//InvalidateRect(hwnd, NULL, TRUE);
			printWindowFunc(hwnd, windowDrawContext, 0);
		SetWorldTransform(windowDrawContext, &identityTransform);
		
		// Sometimes while this update is pending or occuring, part of the
		// window gets invalidated, and it is incorrectly marked as redrawn
		// because of the PrintWindow repaint (even though that went to a
		// different device). To work around this, we invalidate the whole
		// window.
		RedrawWindow(hwnd, NULL, NULL, RDW_ALLCHILDREN|RDW_INVALIDATE|RDW_FRAME);
		
		// Copy onto the combined buffer
		BitBlt(combinedImage, targetRect.left, targetRect.top, targetWidth, targetHeight, windowDrawContext, 0, 0, SRCCOPY);
		
		// Clean up
		/*if(oldRegion)
			SelectObject(windowDrawContext, oldRegion);
		if(region)
			DeleteObject(region);*/
		SelectObject(windowDrawContext, oldWindowBmp);
		DeleteObject(windowDrawBmp);
		DeleteDC(windowDrawContext);
	}
	
	// Record the result
	SelectObject(combinedImage, oldBmp);
	DeleteDC(combinedImage);
	
	return new Snapshot(outputWidth, outputHeight, combinedBmp, getSystemTime());
}

HRGN SnapshotManager::getPreviewRegion(HWND window, const XFORM *transform)
{
	HRGN windowRegion = CreateRectRgn(0,0,0,0);
	
	if(GetWindowRgn(window, windowRegion) == ERROR)
	{
		DeleteObject(windowRegion);
		return NULL;
	}
	
	DWORD regionDataSize = GetRegionData(windowRegion, 0, NULL);
	LPRGNDATA regionData = (LPRGNDATA)HeapAlloc(GetProcessHeap(), 0, regionDataSize);
	GetRegionData(windowRegion, regionDataSize, regionData);
	HeapFree(GetProcessHeap(), 0, regionData);
	
	return ExtCreateRegion(transform, regionDataSize, regionData);
}


/////////////////////////////////////////////////////////////////////////////

bool SnapshotSpec::operator<(const SnapshotSpec &rhs) const
{
	if(task < rhs.task) return true;
	if(task > rhs.task) return false;
	
	if(!vertical && rhs.vertical) return true;
	if(vertical && !rhs.vertical) return false;
	
	if(maxLength < rhs.maxLength) return true;
	if(maxLength > rhs.maxLength) return false;
	
	if(thickness < rhs.thickness) return true;
	if(thickness > rhs.thickness) return false;
	
	return false;
}

bool SnapshotSpec::operator!=(const SnapshotSpec &rhs) const
{
	return task!=rhs.task || vertical!=rhs.vertical
	       || maxLength!=rhs.maxLength || thickness!=rhs.thickness;
}